#include <vector>
#include <string>
#include <thread>

#include "LabController.hpp"
#include "Criminisi.hpp"
#include "Improved.hpp"

using namespace std;

LabController::LabController(std::string benchName)
{
	this->benchName = benchName;
}

LabController::~LabController()
{
	delete this->cf;
}

void LabController::initLab()
{
	this->cf = new yjy::Config("../benchmark-dataset/config.json");

	// retrieve the original image and mask from the benchmark set.
	// mask255 is for display purposes.
	this->getOrigMask(benchName, this->orig, this->mask, this->mask255);

	//back up the original image
	this->orig.copyTo(origBU);

	//get damaged area
	vector<vector<cv::Point>> contours;
	vector<cv::Vec4i> hierarchy;
	findContours(mask, contours, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_SIMPLE);
	double area = contourArea(contours[0]);
	this->damagedPercent = area / orig.rows / orig.cols * 100;

	//calculate window size according to picture size and damagedPercent(mask size),window size must be >= 3
	this->windowSize = this->initWSCalculator(this->damagedPercent, orig.rows * orig.cols);

	// do the actual work

	// display them all

	//output expirmental data

	//save result
}

void LabController::run()
{
	std::thread t1, t2;

	if (this->cf->IsCriminisi())
	{
		t1 = std::thread([](Algorithm *al, int &timeCost, cv::Mat orig, cv::Mat mask, cv::Mat &result, int windowSize) {
			//线程休眠5秒
			//std::this_thread::sleep_for(std::chrono::seconds(5));
			timeCost = al->run(orig, mask, result, windowSize);
			delete al;
		},
						 new Criminisi(), std::ref(this->timeCost_criminisi), this->orig, this->mask, std::ref(this->result_criminisi), this->windowSize);
	}
	if (this->cf->IsImproved())
	{
		t2 = std::thread([](Algorithm *al, int &timeCost, cv::Mat orig, cv::Mat mask, cv::Mat &result, int windowSize) {
			timeCost = al->run(orig, mask, result, windowSize);
			delete al;
		},
						 new Improved(this->cf), std::ref(this->timeCost_improved), this->orig, this->mask, std::ref(this->result_improved), this->windowSize);
	}

	if (this->cf->IsCriminisi())
		t1.join();
	if (this->cf->IsImproved())
		t2.join();
	this->printData();
	this->saveData();
	this->displayResult();

	cv::destroyAllWindows();
}

void LabController::extractMaskFromBench(cv::Mat &maskImg, cv::Mat &mask, unsigned char mask_val, unsigned char red_color_min, unsigned char other_color_max)
{
	mask = cv::Mat::zeros(maskImg.size(), CV_8UC1);

	for (int i = 0; i < maskImg.rows; ++i)
	{
		for (int j = 0; j < maskImg.cols; ++j)
		{
			cv::Vec3b val = maskImg.at<cv::Vec3b>(i, j);
			//Vec3b：蓝-绿-红，红色定为掩码
			if (val[2] >= red_color_min && val[0] <= other_color_max && val[1] <= other_color_max)
			{
				mask.at<unsigned char>(i, j) = mask_val;
			}
		}
	}
}

void LabController::getOrigMask(std::string benchName, cv::Mat &orig, cv::Mat &mask, cv::Mat &mask255)
{
	string BASE_BENCH_ADR = "../benchmark-dataset/";
	string ORIG_APPEND = "-0-original.jpg";
	string MASK_APPEND = "-1-masked.jpg";

	orig = cv::imread(BASE_BENCH_ADR + benchName + ORIG_APPEND);
	cv::Mat maskImg = cv::imread(BASE_BENCH_ADR + benchName + MASK_APPEND);
	extractMaskFromBench(maskImg, mask);
	extractMaskFromBench(maskImg, mask255, 255);
}

void LabController::displayResult()
{
	cv::imshow("orig", origBU);
	if (this->cf->IsCriminisi())
		cv::imshow("result_criminisi", result_criminisi);
	if (this->cf->IsImproved())
		cv::imshow("result_improved", result_improved);
	cv::waitKey(0);
}

void LabController::printData()
{
	cout << "picture size :" << orig.rows << "*" << orig.cols << endl;
	cout << "init window size :" << windowSize << std::endl;
	cout << "damaged percent:" << this->damagedPercent << "%" << endl;
	if (this->cf->IsCriminisi())
		cout << "criminisi patching costs :" << this->timeCost_criminisi << " seconds" << endl;
	if (this->cf->IsImproved())
		cout << "improved patching costs :" << this->timeCost_improved << " seconds" << endl;
}

void LabController::saveData()
{
	if (this->cf->IsCriminisi())
		imwrite("../benchmark-dataset/" + benchName + "-2-result_criminisi.jpg", result_criminisi);
	if (this->cf->IsImproved())
		imwrite("../benchmark-dataset/" + benchName + "-2-result_improved.jpg", result_improved);
}

int LabController::initWSCalculator(float damagedPercent, float pictureSize)
{
	int ws;

	if (!this->cf->variable_init_ws)
		ws = this->cf->init_ws;
	else
	{
		ws = 2 * pictureSize / 1e5 * 0.14f * damagedPercent + 0.5f;
		ws = ws < 5 ? 5 : ws;
	}
	//(0.5f * damagedPercent + 2) * pow(pictureSize / 250000.0f, 0.5f) + 0.5f;
	return ws % 2 ? ws : ws + 1;
}